package com.btl.service.ai.support.rag.embedding;

import com.btl.service.ai.support.rag.RagProperties;
import com.btl.service.ai.support.rag.embedding.models.AllMpNetBaseV2Model;
import com.btl.service.ai.support.rag.embedding.models.BegSmallZhV15Model;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.embedding.onnx.PoolingMode;
import dev.langchain4j.model.embedding.onnx.allminilml6v2.AllMiniLmL6V2EmbeddingModel;
import dev.langchain4j.model.embedding.onnx.bgesmallenv15q.BgeSmallEnV15QuantizedEmbeddingModel;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;

import java.util.concurrent.Executor;
import java.util.concurrent.Executors;

@Configuration(proxyBeanMethods = false)
public class EmbeddingModelProvider {

    private static final Executor executor = Executors.newVirtualThreadPerTaskExecutor();

    @Bean
    public EmbeddingModel bgeSmallEnV15QuantizedEmbeddingModel() {
        return new BgeSmallEnV15QuantizedEmbeddingModel(executor);
    }

    @Primary
    @Bean
    public EmbeddingModel allMiniLmL6V2EmbeddingModel() {
        return new AllMiniLmL6V2EmbeddingModel(executor);
    }

    @Bean
    public EmbeddingModel allMpNetBaseV2Model(RagProperties ragProperties) {
        return new AllMpNetBaseV2Model(
                executor,
                ragProperties.getEmbedding()
                        .getAllMpNetBaseV2()
                        .modelPath(), ragProperties.getEmbedding()
                .getAllMpNetBaseV2()
                .tokenizerPath(),
                PoolingMode.MEAN);
    }

    @Bean
    public EmbeddingModel begSmallZhV15Model(RagProperties ragProperties) {
        return new BegSmallZhV15Model(
                executor,
                ragProperties.getEmbedding()
                        .getBegSmallZhV15()
                        .modelPath(),
                ragProperties.getEmbedding()
                        .getBegSmallZhV15()
                        .tokenizerPath(),
                PoolingMode.MEAN
        );
    }

//    @Bean
//    public EmbeddingModel openAiEmbeddingModel() {
//        return OpenAiOfficialEmbeddingModel.builder()
//                .apiKey("")
//                .modelName("text-embedding-ada-002")
//                .build();
//    }


}
